import socket
import base64
import os
import time
from typing import Tuple, Optional


class UDPClient:
    """UDP file download client"""

    MAX_RETRIES = 5
    INITIAL_TIMEOUT = 1.0  # seconds
    MAX_BLOCK_SIZE = 1000  # bytes

    def __init__(self, host: str, port: int, file_list: str):
        """Initialize the client

        Args:
            host: Server hostname
            port: Server port
            file_list: Path to file containing list of files to download
        """
        self.server_address = (host, port)
        self.file_list = file_list
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.socket.settimeout(self.INITIAL_TIMEOUT)

    def send_and_receive(self, message: str, address: Tuple[str, int],
                         max_retries: int = MAX_RETRIES,
                         initial_timeout: float = INITIAL_TIMEOUT) -> Optional[str]:
        """Send message with timeout and retry mechanism

        Args:
            message: Message to send
            address: Destination address (host, port)
            max_retries: Maximum number of retries
            initial_timeout: Initial timeout in seconds

        Returns:
            Received response message, or None if failed

        Raises:
            Exception: When max retries reached
        """
        current_timeout = initial_timeout
        retries = 0

        while retries < max_retries:
            try:
                self.socket.settimeout(current_timeout)
                start_time = time.time()
                self.socket.sendto(message.encode('utf-8'), address)

                while True:
                    try:
                        response, _ = self.socket.recvfrom(65536)
                        return response.decode('utf-8')
                    except socket.timeout:
                        # Check total timeout period
                        if time.time() - start_time > current_timeout:
                            raise socket.timeout
                        continue

            except socket.timeout:
                retries += 1
                current_timeout *= 2
                print(f"[CLIENT] Timeout, retrying... (attempt {retries}/{max_retries})")
            except Exception as e:
                print(f"[ERROR] Communication error: {e}")
                retries += 1

        raise Exception(f"Max retries ({max_retries}) reached, aborting")

    def _verify_response(self, response: str, expected_start: str) -> bool:
        """Verify server response format

        Args:
            response: Server response
            expected_start: Expected response prefix

        Returns:
            bool: Whether response is valid
        """
        if not response.startswith(expected_start):
            print(f"[ERROR] Invalid server response: {response}")
            return False
        return True

    def _display_progress(self, current: int, total: int):
        """Display download progress

        Args:
            current: Bytes received so far
            total: Total file size in bytes
        """
        percent = (current / total) * 100
        print(f"\rProgress: {current}/{total} bytes ({percent:.1f}%)", end='', flush=True)

    def download_file(self, file_name: str) -> bool:
        """Download a single file

        Args:
            file_name: Filename to download

        Returns:
            bool: Whether download succeeded
        """
        try:
            print(f"\n[CLIENT] Requesting file: {file_name}")

            # Send download request
            response = self.send_and_receive(
                f"DOWNLOAD {file_name}",
                self.server_address
            )

            if not response:
                print("[ERROR] No response from server")
                return False

            print(f"[CLIENT] Server response: {response}")

            if not self._verify_response(response, "OK "):
                return False

            parts = response.split()
            try:
                file_size = int(parts[3])  # OK <filename> SIZE <size> PORT <port>
                port = int(parts[5])
            except (IndexError, ValueError) as e:
                print(f"[ERROR] Failed to parse server response: {e}")
                return False

            print(f"[CLIENT] Downloading {file_name} (size: {file_size} bytes, port: {port})")

            # Prepare local file
            downloaded_path = os.path.join("client_files", file_name)
            os.makedirs(os.path.dirname(downloaded_path), exist_ok=True)

            with open(downloaded_path, 'wb') as file:
                bytes_received = 0
                data_port = (self.server_address[0], port)

                while bytes_received < file_size:
                    start = bytes_received
                    end = min(start + self.MAX_BLOCK_SIZE - 1, file_size - 1)

                    # Request data block
                    request = f"FILE {file_name} GET START {start} END {end}"
                    response = self.send_and_receive(request, data_port)

                    if not response:
                        print("[ERROR] No response for data block")
                        return False

                    if not self._verify_response(response, f"FILE {file_name} OK"):
                        return False

                    # Parse data
                    try:
                        data_index = response.find("DATA") + 5
                        if data_index <= 5:
                            raise ValueError("DATA marker not found")

                        base64_data = response[data_index:].strip()
                        file_data = base64.b64decode(base64_data)
                        file.seek(start)
                        file.write(file_data)
                        bytes_received += len(file_data)
                        self._display_progress(bytes_received, file_size)
                    except Exception as e:
                        print(f"\n[ERROR] Data decoding error: {e}")
                        return False

                # Send close request and verify response
                close_msg = f"FILE {file_name} CLOSE"
                response = self.send_and_receive(close_msg, data_port)

                if not response or not self._verify_response(response, f"FILE {file_name} CLOSE_OK"):
                    print("[ERROR] File transfer did not complete properly")
                    return False

            print(f"\n[CLIENT] {file_name} downloaded successfully")
            return True

        except Exception as e:
            print(f"\n[ERROR] Error downloading {file_name}: {e}")
            return False

    def start(self):
        """Start the client download process"""
        try:
            if not os.path.exists(self.file_list):
                raise FileNotFoundError(f"File list {self.file_list} not found")

            with open(self.file_list, 'r', encoding='utf-8') as f:
                files = [line.strip() for line in f if line.strip()]

            if not files:
                print("[CLIENT] File list is empty")
                return

            print(f"[CLIENT] Starting download of {len(files)} files...")
            for file_name in files:
                if not self.download_file(file_name):
                    print(f"[WARNING] Failed to download {file_name}, continuing with next file")

        except Exception as e:
            print(f"[ERROR] Client error: {e}")
        finally:
            self.socket.close()
            print("[CLIENT] Client shutdown")


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 4:
        print("Usage: python UDPClient.py <host> <port> <file_list>")
        sys.exit(1)

    client = UDPClient(sys.argv[1], int(sys.argv[2]), sys.argv[3])
    client.start()